# -*- coding: UTF-8 -*-
import os
from PIL import Image
import torch
from torch.autograd import Variable
from torchvision.transforms import ToTensor, ToPILImage
from util.RRDBNet import RRDBNet
from util.runtimeError import runtimeError

'''
完成超分辨功能：
1、expandeImage完成图片4倍放大
2、restore_orin_size实现在保证图片原尺寸不变下，重建更多细节信息
'''
global model
model = RRDBNet(3, 3, 64, 23, gc=32)


def expandImage(filepaths, rootpath):
    try:
        # 获得图片路径
        pathlist = filepaths.split('\n\n')
        pathlist[-1] = pathlist[-1].rstrip('\n')
        for i in range(len(pathlist)):
            pathlist[i] = pathlist[i].split('->')[1]
        for index, path in enumerate(pathlist, 1):
            image = Image.open(path)
            with torch.no_grad():
                image = Variable(ToTensor()(image)).unsqueeze(0)
            if torch.cuda.is_available():
                image = image.cuda()
                model.cuda()
                model.load_state_dict(torch.load(rootpath + '/data/RRDB_ESRGAN_x4.pth'))
            else:
                image = image.cpu()
                model.cpu()
                model.load_state_dict(torch.load(rootpath + '/data/RRDB_ESRGAN_x4.pth', map_location=lambda storage, loc: storage))
            out = model(image)
            out = torch.clamp(out, 0.0, 1.0)
            out_image = ToPILImage()(out[0].data.cpu())
            if not os.path.exists(rootpath + '/output/SR/upsample/'):
                os.makedirs(rootpath + '/output/SR/upsample/')
            out_image.save(rootpath + '/output/SR/upsample/' + str(index) + '.png')
    except Exception as e:
        if e.args[0].find("CUDA out of memory.") != -1:
            raise runtimeError("处理图片" + str(index) + " 遇到问题。显存不够。具体：" + e.args[0])
        else:
            raise runtimeError("处理图片" + str(index) + " 遇到问题。具体：" + e.args[0])


def restore_orin_size(filepaths, rootpath):
    try:
        pathlist = filepaths.split('\n\n')
        pathlist[-1] = pathlist[-1].rstrip('\n')
        for i in range(len(pathlist)):
            pathlist[i] = pathlist[i].split('->')[1]
        for index, path in enumerate(pathlist, 1):
            # 目前只做4倍放大，所以需要适应性裁剪图片
            image = adoptsize(Image.open(path))
            image_downsample = image.resize((image.width // 4, image.height // 4), resample=Image.BICUBIC)
            with torch.no_grad():
                image_downsample = Variable(ToTensor()(image_downsample)).unsqueeze(0)
            if torch.cuda.is_available():
                image_downsample = image_downsample.cuda()
                model.cuda()
                model.load_state_dict(torch.load(rootpath + '/data/RRDB_ESRGAN_x4.pth'))
            else:
                image_downsample = image_downsample.cpu()
                model.cpu()
                model.load_state_dict(torch.load(rootpath + '/data/RRDB_ESRGAN_x4.pth', map_location=lambda storage, loc: storage))
            out = model(image_downsample)
            out = torch.clamp(out, 0.0, 1.0)
            out_image = ToPILImage()(out[0].data.cpu())
            if not os.path.exists(rootpath + '/output/SR/restore/'):
                os.makedirs(rootpath + '/output/SR/restore/')
            out_image.save(rootpath + '/output/SR/restore/' + str(index) + '.png')
    except Exception as e:
        if e.args[0].find('CUDA out of memory') != -1:
            raise runtimeError("处理图片" + str(index) + " 遇到问题。显存不够。具体：" + e.args[0])
        else:
            raise runtimeError("处理图片" + str(index) + " 遇到问题。具体：" + e.args[0])


def adoptsize(img):
    width, height = img.size
    adoptw = width - width % 4
    adopth = height - height % 4
    outimg = img.resize((adoptw, adopth), resample=Image.BICUBIC)
    return outimg
